'''
Created on May 1, 2010

@author: Guan Gui, Will Zhang

'''

from collections import defaultdict
from alignment import Alignment
import sys

class EMIBMModel1(object):
    '''
    This class contains implementations of the Expectation Maximization
    algorithm for IBM Model 1. The algorithm runs upon a sentence-aligned 
    parallel corpus and generates word alignments in aligned sentence pairs.
    
    The process is divided into 2 main stages.
    Stage 1: Studies word-to-word translation probabilities by collecting
    evidence of a English word been the translation of a foreign word from
    the parallel corpus.
    
    Stage 2: Based on the translation probabilities from Stage 1, generates 
    word alignments for aligned sentence pairs.
    '''

    
    NULL_TOKEN = "#NULL"
    # This field is internal to the algorithm implementation.

    class TableCell(object):
        '''
        Instances of this class are used as entries of the t(e|f)
        translation probability table.
        
        The translation probability is in instance variable "t"
        
        Other fields and all methods of this class are internal to the
        outer class implementation, they should not be accessed or invoked
        outside this module.
        '''
        init_t = 0

        def __init__(self):
            self.t = self.init_t
            self.count = 0

        def clearCount(self):
            self.count = 0

        def updateT(self, total_f):
            t_history = self.t
            self.t = self.count / total_f
            return abs(t_history - self.t)

    def __init__(self, aligned_sents, convergent_threshold = 1 * 1e-2, debug =
            False):
        '''
        Initialize a new C{EMIBMModel1}.

        @param aligned_sents: The parallel text corpus.Iteratable containing 
            AlignedSent instances of aligned sentence pairs from the corpus.
        @type aligned_sents: C{list} of L{AlignedSent} objects
        @param convergent_threshold: The threshold value of convergence. An 
            entry is considered converged if the delta from old_t to new_t
            is less than this value. The algorithm terminates when all entries
            are converged. This parameter is optional, default is 0.01
        @type convergent_threshold: C{int}
        @param debug: debug flag used to suppress the progress report output in
        order to do doctest
        @type debug: C{bool}
        '''
        self.aligned_sents = aligned_sents
        self.convergent_threshold = convergent_threshold
        # Dictionary of translation probabilities.
        self.e_given_f = None
        self.debug = debug
        self.train()

    def t_e_given_f(self, e_word, f_word):
        '''
        Accessor method for translation probabilities t(e|f)

        @param e_word: Look up English Word
        @type e_word: C{str}
        @param f_word: Foreign Word translating to. For NULL token, use the
            constant "EMIBMModel1.NULL_TOKEN"
        @type f_word: C{str}
        @return: The translation probability of e_word to f_word
        @rtype: C{int}
        '''
        return self.e_given_f[(e_word, f_word)].t

    def train(self):
        '''
        The train() function implements Expectation Maximization training
        stage that learns word-to-word translation probabilities.        
        '''
        # generate word sets from text corpus
        e = set()
        f = set()
        # this is necessary since we don't want to modify the original sets
        for aligned_sent in self.aligned_sents:
            for e_w in aligned_sent.words:
                e.add(e_w)
            for f_w in aligned_sent.mots:
                f.add(f_w)
        # add the NULL token to the foreign word set.
        f.add(self.NULL_TOKEN)
        
        entry_count = len(e) * len(f)
        
        # Initialization
        self.e_given_f = defaultdict(self.TableCell)
        total = defaultdict(float)
        s_total = defaultdict(float)

        # Initialize t(e|f) uniformly
        self.TableCell.init_t = float(1) / len(e)

        globally_converged = False
        iteration_count = 0

        while not globally_converged:
            iteration_count += 1

            for aligned_sent in self.aligned_sents:
                for e_w in aligned_sent.words:
                    s_total[e_w] += self.e_given_f[(e_w, self.NULL_TOKEN)].t
                    for f_w in aligned_sent.mots:
                        s_total[e_w] += self.e_given_f[(e_w, f_w)].t

                # collect counts
                for e_w in aligned_sent.words:
                    cur_e_given_f = self.e_given_f[(e_w, self.NULL_TOKEN)]
                    tmp = cur_e_given_f.t / s_total[e_w]
                    cur_e_given_f.count += tmp
                    total[self.NULL_TOKEN] += tmp
                    for f_w in aligned_sent.mots:
                        cur_e_given_f = self.e_given_f[(e_w, f_w)]
                        tmp = cur_e_given_f.t / s_total[e_w]
                        cur_e_given_f.count += tmp
                        total[f_w] += tmp

            globally_converged = True
            conv_count = 0
            for f_w in f:
                for e_w in e:
                    cur_e_given_f = self.e_given_f[(e_w, f_w)]
                    diff = cur_e_given_f.updateT(total[f_w])
                    entry_converged = diff <= self.convergent_threshold
                    globally_converged = globally_converged and entry_converged
                    conv_count += entry_converged
                    cur_e_given_f.clearCount()

                total[f_w] = 0
            if not self.debug:
                # Interactive status report of the progress to standard out
                print "%d/%d (%5.4f%%) entries have converged with \
threshold %.2e\r" % (conv_count, entry_count, conv_count * 100.0 / entry_count,
                            self.convergent_threshold),
                sys.stdout.flush() # this must be flushed to see the latest 
                                   # result
        if not self.debug:
            print
        print "Finished training after %d iterations." % (iteration_count)

    def dump_t_e_given_f(self, dump_threshold = 0.75):
        '''
        Dump learnt t(e|f)s that have probabilities >= dump_threshold

        @param dump_threshold: Threshold to dump t(e|f) with. Default to 0.75.
        @return: the string with each 't(e|f) = prob' entry in a line
        @rtype: C{str}
        '''
        s = ''
        # Selective t(e|f) Dump, for demonstration only
        for k, v in self.e_given_f.iteritems():
            if v.t > dump_threshold:
                s += "t(%s|%s) = %.16f\n" % (k[0], k[1], v.t)
        return s

    def yield_alignments(self):
        '''
        The yield_alignments() function implements the alignment learning
        stage. The yielded alignments are stored back to those AlignedSent objects.
        '''
        # Alignment Learning from t(e|f)
        for aligned_sent in self.aligned_sents:

            alignment = [];

            # for every English word
            for j, e_w in enumerate(aligned_sent.words):
                # find the French word that gives maximized t(e|f)
                # NULL_TOKEN is the initial candidate
                f_index, t_max = -1, self.e_given_f[(e_w, self.NULL_TOKEN)].t
                for i, f_w in enumerate(aligned_sent.mots):
                    t = self.e_given_f[(e_w, f_w)].t
                    if t > t_max:
                        f_index, t_max = i, t
                
                # only output alignment with non-NULL mapping
                if f_index != -1:
                    alignment.append((j,f_index))

            # substitute the alignment of AlignedSent with the yielded one
            aligned_sent.alignment = Alignment(alignment)
